from torch.utils.data import Dataset
import random
import numpy as np
from PIL import Image
import torchvision
from datasets.Asymmetric_Noise import noisify_cifar100_asymmetric
import torch
from utils.mypath import Path

from datasets.imagenet32 import Imagenet32
from datasets.utils_imbalance import get_imb_data


class cifar_dataset(Dataset):
    def __init__(self, dataset, r_ood=0.2, r_id=0.2, r_imb=0.1, transform=None, mode=None, seed=0, asym=False,
                 ood=None, pred=None, probability=None,
                 transform_st=None):
        self.mode = mode
        self.probability = probability
        self.transform = transform
        self.transform_st = transform_st
        if dataset == 10:
            dataset = 'cifar10'
        elif dataset == 100:
            dataset = 'cifar100'
        if mode == 'test':
            if dataset == 'cifar10':
                cifar = CIFAR10(train=False, transform=transform)
            elif dataset == 'cifar100':
                cifar = CIFAR100(train=False, transform=transform)
            self.test_data = cifar.data
            self.test_label = cifar.targets
            return
        else:
            if dataset == 'cifar10':
                cifar = CIFAR10(train=True, transform=transform, asym=asym, r_ood=r_ood, r_id=r_id, r_imb=r_imb, seed=seed)
            elif dataset == 'cifar100':
                cifar = CIFAR100(train=True, transform=transform, asym=asym, r_ood=r_ood, r_id=r_id, r_imb=r_imb, seed=seed)
            train_data = cifar.data
            noise_label = cifar.targets

        if self.mode == 'ood':
            ood_idx = ood.nonzero()[0]
            self.train_data = train_data[ood_idx]
            self.noise_label = [noise_label[i] for i in ood_idx]
            print("%s data has a size of %d" % (self.mode, len(self.noise_label)))

        if self.mode == 'all' or self.mode == 'all_sup':
            self.train_data = train_data
            self.noise_label = noise_label
        elif self.mode == 'all-ood':
            ind_idx = ood.zero()[0]
            self.train_data = train_data[ind_idx]
            self.noise_label = [noise_label[i] for i in ind_idx]
            print("%s data has a size of %d" % (self.mode, len(self.noise_label)))
        else:
            if self.mode == "labeled":
                pred_idx = (pred * ood).nonzero()[0]
                self.probability = [probability[i] for i in pred_idx]

            elif self.mode == "unlabeled":
                pred_idx = ((1 - pred) * ood).nonzero()[0]

            self.train_data = train_data[pred_idx]
            self.noise_label = [noise_label[i] for i in pred_idx]
            print("%s data has a size of %d" % (self.mode, len(self.noise_label)))

    def __getitem__(self, index):
        if self.mode == 'labeled':
            img, target, prob = self.train_data[index], self.noise_label[index], self.probability[index]
            target = torch.tensor(target).long()
            img = Image.fromarray(img)
            if self.transform_st:
                img1 = self.transform(img)
                img2 = self.transform(img)
                img3 = self.transform_st(img)
                img4 = self.transform_st(img)
                return img1, img2, img3, img4, target, prob
            else:
                img1 = self.transform(img)
                img2 = self.transform(img)
                return img1, img2, target, prob
        elif self.mode == 'unlabeled':
            img = self.train_data[index]
            img = Image.fromarray(img)
            if self.transform_st:
                img1 = self.transform(img)
                img2 = self.transform(img)
                img3 = self.transform_st(img)
                img4 = self.transform_st(img)
                return img1, img2, img3, img4
            else:
                img1 = self.transform(img)
                img2 = self.transform(img)
                return img1, img2
        elif self.mode == 'all':
            img, target = self.train_data[index], self.noise_label[index]
            img = Image.fromarray(img)
            img1, img2 = self.transform(img)
            return (img1, img2), target, index
        elif self.mode == 'all_sup':
            img, target = self.train_data[index], self.noise_label[index]
            img = Image.fromarray(img)
            img1 = self.transform(img)
            return img1, target, index
        elif self.mode == 'test':
            img, target = self.test_data[index], self.test_label[index]
            img = Image.fromarray(img)
            img = self.transform(img)
            return img, target

    def __len__(self):
        if self.mode != 'test':
            return len(self.train_data)
        else:
            return len(self.test_data)


class CIFAR10(torchvision.datasets.CIFAR10):
    nb_classes = 10

    def __init__(self, root='~/data', train=True, transform=None,
                 r_ood=0.2, r_id=0.2, r_imb=0.1, seed=0, asym=False):

        super().__init__(root, train=train, transform=transform)
        if train is False:
            return
        if r_imb > 0.:
            self.data, self.targets = get_imb_data(self.data, self.targets,
                                                   self.nb_classes, r_imb, seed)
            print('built imbalanced dataset')
        np.random.seed(seed)
        if r_ood > 0.:
            ids_ood = [i for i in range(len(self.targets)) if np.random.random() < r_ood]
            imagenet32 = Imagenet32(root='~/data/imagenet32', train=True)
            img_ood = imagenet32.data[np.random.permutation(range(len(imagenet32)))[:len(ids_ood)]]
            self.ids_ood = ids_ood
            self.data[ids_ood] = img_ood
            print('Mixing in OOD noise')

        if r_id > 0.:
            ids_not_ood = [i for i in range(len(self.targets)) if i not in ids_ood]
            ids_id = [i for i in ids_not_ood if np.random.random() < (r_id / (1 - r_ood))]
            if asym:
                if self.nb_classes == 10:
                    transition = {0: 0, 2: 0, 4: 7, 7: 7, 1: 1, 9: 1, 3: 5, 5: 3, 6: 6, 8: 8}
                    for i, t in enumerate(self.targets):
                        if i in ids_id:
                            self.targets[i] = transition[t]
                else:
                    self.targets, _ = noisify_cifar100_asymmetric(self.targets, r_id, seed)
                print('Mixing in ID asym noise')
            else:
                for i, t in enumerate(self.targets):
                    if i in ids_id:
                        self.targets[i] = int(np.random.random() * self.nb_classes)
                print('Mixing in ID noise')
            self.ids_id = ids_id


class CIFAR100(CIFAR10):
    base_folder = "cifar-100-python"
    url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
    filename = "cifar-100-python.tar.gz"
    tgz_md5 = "eb9058c3a382ffc7106e4002c42a8d85"
    train_list = [
        ["train", "16019d7e3df5f24257cddd939b257f8d"],
    ]

    test_list = [
        ["test", "f0ef6b0ae62326f3e7ffdfab6717acfc"],
    ]
    meta = {
        "filename": "meta",
        "key": "fine_label_names",
        "md5": "7973b15100ade9c7d40fb424638fde48",
    }

    nb_classes = 100
